# -*- coding: utf-8 -*-
"""
Created on Sat Aug 17 19:09:28 2024

@author: TAO

根据图片分析程序， 计算小球的瞬时位移，瞬时速度， 瞬时加速
并且以图表的形式显示处来
"""

import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.signal import savgol_filter


def read_excel_file(excel_file_path,scale_factor):
    df = pd.read_excel(excel_file_path, sheet_name='Sheet1')

    # 过滤掉以 "T_" 开头的行,T_开头的图片对应的是模板图片
    df = df[~df['ImageName'].str.startswith('T_')]

    #采样间隔，由于每帧采样的间距较近，稍有偏差就导致结果跳动巨大
    sampling_interval = 1 #表示采样步长，隔sampling_interval行采样一次，=3表示步长是3，隔两行采样一次
    # 提取BallX和BallY两列
    ball_coordinates = df[['BallX', 'BallY']].iloc[::sampling_interval]

    # 转换为列表的列表形式
    ball_coordinates_lists = ball_coordinates.values.tolist()
    
   
    # 矫正物理量纲，opencv 获取的是图片的像素坐标，这里要换算成物理量纲    
    # 创建一个空列表来存储归一化后的坐标
    normalized_ball_coordinates_lists = []
    
    # 对应 `ball_coordinates_lists` 中的每一行
    for row in ball_coordinates_lists:
        # 归一化每一行中的每个元素
        row[0] = row[0] / scale_factor
        row[1] = row[1] / scale_factor
        # 将归一化后的新行添加到最终列表中
        normalized_ball_coordinates_lists.append(row)
    
    return normalized_ball_coordinates_lists

def plot_data(df, title, xlabel, ylabel, y_min=None, y_max=None):
    ax = df.plot(x='x', y='y', kind='line', title=title, grid=True, xlabel=xlabel, ylabel=ylabel)
    if y_min is not None and y_max is not None:
        ax.set_ylim([y_min, y_max])
    plt.show()

def motion_calc(excel_file_path,scale_factor):

    # 打开excel文件，返回小球坐标(x,y)的列表，每行的时间间隔是0.016秒
    # 这里返回的列表，已经是归一化的，坐标的单位都是米
    ball_coordinates_lists = read_excel_file(excel_file_path,scale_factor)
    
    # 平滑参数
    window_length = 5  # 窗口大小
    polyorder = 2      # 多项式阶数
    #每帧的帧间隔是0.016秒
    dt = 1.0/60.0
    # 遍历列表中的每个坐标对
  
    t_for_x = [] #X 轴的位移时间序列，从t0时刻开始
    t_for_v = [] #X 轴的速度时间序列，从t1时刻开始
    t_for_a = [] #X 轴的加速度时间序列，从t2时刻开始
    
    x_x = []   #每时刻横向位移
    y_y = []   #每时刻纵向位移
    x_y = []   #XY合力方向的位移
    v_x = []   #X 轴方向上每个点的瞬时速度
    v_y = []   #Y 轴方向上每个点的瞬时速度
    v_xy= []   #XY合力方向的矢量速度
    a_x = []   #X 轴方向上每个点的瞬时加速度集合 水平加速度 ax = Δvx / Δt
    a_y = []   #Y 轴方向上每个点的瞬时加速度集合 垂直加速度 ay = Δvy / Δt
    a_xy= []   #XY合力方向的矢量加速度
    

    #提取每一帧,X，Y的坐标
    for i in range(0, len(ball_coordinates_lists)):
        coord_curr = ball_coordinates_lists[i]   #当前点
        print(f"坐标对 {i}: BallX = {coord_curr[0]}, BallY = {coord_curr[1]}")
        t_for_x.append(i)   #当前时刻的t 帧数                
        x_x.append(coord_curr[0]) #X如果从右往左抛，就是负值
        y_y.append(coord_curr[1]) #Y如果从上往下运动，就是负值
        if i == 0:
            x_y.append(np.sqrt(x_x[i]**2 + y_y[i]**2)) #未平滑的数据计算第0个元素
    #end of for        
    # 使用平滑滤波器处理数据，消除数据抖动,这里平滑位移
    x_x_smoothed = savgol_filter(x_x, window_length, polyorder)
    y_y_smoothed = savgol_filter(y_y, window_length, polyorder)
    
    x_x = x_x_smoothed
    y_y = y_y_smoothed
    
    for j in range(1, len(x_x)):
        x_y.append(np.sqrt(x_x[j]**2 + y_y[j]**2))   #求出t_i 时刻XY合力方向位移,用的是平滑后的数据
        
        dx = x_x[j] - x_x[j-1]          #求出t_j 时刻X的Δ位移
        dy = y_y[j] - y_y[j-1]          #求出t_j-1 时刻Y的Δ位移
        
        vx = abs(dx) / dt               #求出t_j 时刻横向的瞬时速度
        vy = abs(dy) / dt               #求出t_j 时刻纵的瞬时速度
        vxy = np.sqrt(vx**2 + vy**2)    #求出t_j 时刻小球的合力方向瞬时速度

        t_for_v.append(j)               #平均速度对应的是t1时刻
        v_x.append(vx)
        v_y.append(vy)
        v_xy.append(vxy)
      
     #end of for
     
     # 使用平滑滤波器处理数据，消除数据抖动，这里平滑速度
    v_x = savgol_filter(v_x, window_length, polyorder)
    v_y = savgol_filter(v_y, window_length, polyorder)
     # t1 时刻，在这一时刻，我们知道小球的新位置 (X, Y)，但仍然缺乏足够的信息来计算加速度，
     # 因为我们还需要第三个时间点的数据来确定速度的变化，所以t1时刻的加速度还是无意义  
    for k in range(1, len(v_x)):
        t_for_a.append(k)
        ax = abs(v_x[k]-v_x[k-1])/dt   #水平加速度 ax = Δvx / Δt
        ay = abs(v_y[k]-v_y[k-1])/dt   #垂直加速度 ay = Δvy / Δt
        axy= (v_xy[k]-v_xy[k-1])/dt #合力加速度 axy = Δvxy / Δt
        
        a_x.append(ax)
        a_y.append(ay)
        a_xy.append(axy)
     #end of for
     #**画图**
    xx_t = pd.DataFrame({'y': x_x, 'x': t_for_x})  #X横向位移图
    yy_t = pd.DataFrame({'y': y_y, 'x': t_for_x})  #Y纵向位移图
    xy_t = pd.DataFrame({'y': x_y, 'x': t_for_x})  #XY合位移图
    
    vx_t = pd.DataFrame({'y': v_x, 'x': t_for_v})  #X横向速度图
    vy_t = pd.DataFrame({'y': v_y, 'x': t_for_v})  #Y纵向速度图
    vxy_t =  pd.DataFrame({'y': v_xy, 'x': t_for_v}) #XY合力速度图
    
    ax_t = pd.DataFrame({'y': a_x, 'x': t_for_a})   #X 加速度图
    ay_t = pd.DataFrame({'y': a_y, 'x': t_for_a})    #Y 加速度图
    axy_t =  pd.DataFrame({'y': a_xy, 'x': t_for_a}) #XY 加速度图
    
    
    # 使用DataFrame的plot方法绘制折线图
    plot_data(xx_t, 'x_x vs Time', 'Time (Frame)', 'X Displacement (m)')
    plot_data(yy_t, 'y_y vs Time', 'Time (Frame)', 'Y Displacement (m)')
    plot_data(xy_t, 'x_y vs Time', 'Time (Frame)', 'Combined Displacement (m)')
    
    # 设置速度图
    plot_data(vx_t, 'vx vs Time', 'Time (Frame)', 'X Velocity (m/s)',0,5)
    plot_data(vy_t, 'vy vs Time', 'Time (Frame)', 'Y Velocity (m/s)',0,5)
    plot_data(vxy_t, 'vxy vs Time', 'Time (Frame)', 'Combined Velocity (m/s)',0,5)
    
    # 设置加速度图,加速度给定范围，后显示就正确了
    plot_data(ax_t, 'ax vs Time', 'Time (Frame)', 'X Acceleration (m/s²)',0,10)
    plot_data(ay_t, 'ay vs Time', 'Time (Frame)', 'Y Acceleration (m/s²)',0,10)
    plot_data(axy_t, 'axy vs Time', 'Time (Frame)', 'Combined Acceleration (m/s²)',0,10)
   
#------------------------------------------------------------------------------
# 主函数从这里开始
if __name__ == "__main__":
    #需要处理的畸变的图片路径，包含若干张图片,坐标文件也会保存在这个目录下
    #这个目录已经被ImageParser.py 成功处理过，已经产生了坐标文件
    input_image_directory= 'E:/oball/Bin/20240908_141152'
    # 指定输出的excel文件,保存在图片目录下
    excel_file_path = os.path.join(input_image_directory, 'coordinates.xlsx')
    
    #矫正系数,屏幕上每1155像素对应0.124米，因此1米对应9240像素
    scale_factor = 1240.0
    
    motion_calc(excel_file_path,scale_factor)



